import os
import numpy as np
from datetime import datetime
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
from ultralytics.utils.plotting import colors


def soft_nms(dets, scores, sigma=0.5, thresh=0.3, method=1):
    """Soft-NMS实现（纯NumPy）"""
    N = dets.shape[0]
    indexes = np.arange(N)

    for i in range(N):
        # 找到当前最高分框
        max_pos = i + np.argmax(scores[i:])
        if max_pos != i:
            # 交换位置
            dets[[i, max_pos]] = dets[[max_pos, i]]
            scores[i], scores[max_pos] = scores[max_pos], scores[i]
            indexes[i], indexes[max_pos] = indexes[max_pos], indexes[i]

        # 抑制处理
        for pos in range(i + 1, N):
            # 计算IoU
            x1, y1, x2, y2 = dets[i]
            x1_p, y1_p, x2_p, y2_p = dets[pos]
            inter = max(0, min(x2, x2_p) - max(x1, x1_p)) * max(0, min(y2, y2_p) - max(y1, y1_p))
            iou = inter / ((x2 - x1) * (y2 - y1) + (x2_p - x1_p) * (y2_p - y1_p) - inter + 1e-10)

            # 高斯衰减
            weight = np.exp(-(iou ** 2) / sigma)
            scores[pos] *= weight

        # 移除低分框
        mask = scores >= thresh
        dets = dets[mask]
        scores = scores[mask]
        indexes = indexes[mask]
        N = len(dets)

    return indexes


def picture_flag(model, picture_path, result_image_dir):
    os.makedirs(result_image_dir, exist_ok=True)

    # 1. 用Pillow读取图片（RGB格式）
    img = Image.open(picture_path).convert("RGB")
    img_np = np.array(img)  # 转为NumPy数组供YOLO处理

    # 2. 使用模型预测（禁用内置NMS）
    results = model.predict(img_np, nms=False)

    # 3. 提取原始检测结果
    boxes = results[0].boxes
    dets = boxes.xyxy.cpu().numpy()
    scores = boxes.conf.cpu().numpy()
    cls = boxes.cls.cpu().numpy()

    # 4. 应用Soft-NMS
    if len(dets) > 0:
        keep = soft_nms(dets, scores, sigma=0.5, thresh=0.3)
        dets = dets[keep]
        scores = scores[keep]
        cls = cls[keep]

    # 5. 用Pillow绘制结果
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("arial.ttf", size=20)
    except:
        font = ImageFont.load_default()

    for box, score, class_id in zip(dets, scores, cls):
        x1, y1, x2, y2 = map(int, box)
        class_id = int(class_id)

        # 获取类别颜色（RGB格式）
        color = colors(class_id)

        # 绘制矩形框
        draw.rectangle([x1, y1, x2, y2], outline=color, width=2)

        # 绘制标签
        label = f"{results[0].names[class_id]} {score:.2f}"
        draw.text((x1, y1 - 20), label, fill=color, font=font)

    # 6. 保存RGB图片
    current_time = datetime.now().strftime("%m%d%H%M%S")
    output_path = os.path.join(result_image_dir, f"{current_time}_0.jpg")
    img.save(output_path, quality=95)

    # 7. 返回RGB图像（Pillow Image对象）
    return img


def main():
    # 加载训练好的模型
    model = YOLO("best.pt")

    # 图片路径
    image_path = "img_4.png"

    # 结果保存目录
    result_dir = os.path.join("picture_inference", "picture_flag")

    # 进行检测
    result_image = picture_flag(model, image_path, result_dir)

    print(f"检测完成，结果已保存到: {result_dir}")


if __name__ == "__main__":
    main()